// https://leetcode.cn/problems/binary-tree-pruning/
// Created by ade on 2022/7/21.
// 给你二叉树的根结点 root ，此外树的每个结点的值要么是 0 ，要么是 1 。
//
// 返回移除了所有不包含 1 的子树的原二叉树。
//
// 节点 node 的子树为 node 本身加上所有 node 的后代。

#include <iostream>
#include <vector>


using namespace std;

struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;

    TreeNode() : val(0), left(nullptr), right(nullptr) {}

    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}

    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
};

class Solution {
public:
    TreeNode *pruneTree(TreeNode *root) {
        if (!root) return nullptr;
        root->left = pruneTree(root->left);
        root->right = pruneTree(root->right);
        if (root->val == 0 && !root->left && !root->right) return nullptr;
        return root;
    }


    TreeNode *init() {
        TreeNode *head1 = new TreeNode(1);
        TreeNode *head2 = new TreeNode(0);
        TreeNode *head3 = new TreeNode(1);
        TreeNode *head4 = new TreeNode(0);
//        TreeNode *head4 = new TreeNode(3);
//        TreeNode *head5 = new TreeNode(3);
//        TreeNode *head6 = new TreeNode(4);
//        TreeNode *head7 = new TreeNode(4);
        head1->right = head2;
//        head1->right = head3;
        head2->right = head3;
        head2->left = head4;
//
//        head3->left = nullptr;
//        head3->right = nullptr;
//        head4->left = head6;
//        head4->right = head6;
        return head1;
    }

    void show(TreeNode *node) {
        if (!node) return;
        cout << node->val << ",";
        show(node->left);
        show(node->right);
    }
};

int main() {
    Solution so;
    TreeNode *head = so.init();
    so.pruneTree(head);
    so.show(head);
//    cout << head->val << endl;
    return 0;
}